import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim

# 模拟数据
# 假设有3个类别
num_classes = 2
# 假设有一个输入特征张量，batch_size = 5
# inputs = torch.randn(5, requires_grad=True)
inputs = torch.tensor([[0], [1]]).float()
# 假设目标标签是一个包含类别索引的张量
targets = torch.tensor([1, 0])

# 定义模型（定义一个全连接层）
linear = nn.Linear(in_features=1, out_features=num_classes)
# 定义优化器
optimizer = optim.SGD(linear.parameters(), lr=0.01)  # 学习率设置为0.01
# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 训练模型
for epoch in range(5):
    # 前向传播（进行线性变换）
    logits = linear(inputs)
    print(f"logits={logits}")
    print(f"linear.weight={linear.weight}")
    print(f"linear.bias={linear.bias}")
    # 计算损失
    loss = criterion(F.log_softmax(logits, dim=1), targets)
    # 反向传播和优化
    optimizer.zero_grad()  # 清空梯度
    loss.backward()  # 反向传播
    optimizer.step()  # 更新权重

if __name__ == '__main__':
    print("end")
